import os
import re
from typing import List, Dict
import PyPDF2
from docx import Document

class DocumentProcessor:
    """文档处理器，支持多种文件格式"""
    
    def __init__(self):
        self.chunk_size = 500  # 每个文本块的最大字符数
        self.chunk_overlap = 50  # 文本块之间的重叠字符数
    
    def process_file(self, file_path: str) -> List[Dict[str, str]]:
        """处理单个文件"""
        file_extension = os.path.splitext(file_path)[1].lower()
        
        if file_extension == '.txt':
            return self.process_txt(file_path)
        elif file_extension == '.pdf':
            return self.process_pdf(file_path)
        elif file_extension == '.md':
            return self.process_md(file_path)
        elif file_extension == '.docx':
            return self.process_docx(file_path)
        else:
            print(f"不支持的文件格式: {file_extension}")
            return []
    
    def process_txt(self, file_path: str) -> List[Dict[str, str]]:
        """处理文本文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            return self.chunk_text(content)
        except Exception as e:
            print(f"处理文本文件失败: {e}")
            return []
    
    def process_pdf(self, file_path: str) -> List[Dict[str, str]]:
        """处理PDF文件"""
        try:
            text = ""
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                for page in pdf_reader.pages:
                    text += page.extract_text()
            return self.chunk_text(text)
        except Exception as e:
            print(f"处理PDF文件失败: {e}")
            return []
    
    def process_md(self, file_path: str) -> List[Dict[str, str]]:
        """处理Markdown文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            return self.chunk_text(content)
        except Exception as e:
            print(f"处理Markdown文件失败: {e}")
            return []
    
    def process_docx(self, file_path: str) -> List[Dict[str, str]]:
        """处理Word文档"""
        try:
            doc = Document(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return self.chunk_text(text)
        except Exception as e:
            print(f"处理Word文档失败: {e}")
            return []
    
    def chunk_text(self, text: str) -> List[Dict[str, str]]:
        """将文本分割成块"""
        # 清理文本
        text = re.sub(r'\s+', ' ', text.strip())
        
        if not text:
            return []
        
        chunks = []
        start = 0
        
        while start < len(text):
            # 计算结束位置
            end = start + self.chunk_size
            
            # 如果超出文本长度，调整到文本末尾
            if end >= len(text):
                chunk = text[start:]
            else:
                # 尝试在句子边界处分割
                chunk = text[start:end]
                # 查找最后一个句号、问号或感叹号
                last_sentence_end = max(
                    chunk.rfind('。'),
                    chunk.rfind('？'),
                    chunk.rfind('！'),
                    chunk.rfind('.'),
                    chunk.rfind('?'),
                    chunk.rfind('!')
                )
                
                if last_sentence_end > self.chunk_size * 0.7:  # 如果找到了合适的分割点
                    chunk = chunk[:last_sentence_end + 1]
                    end = start + len(chunk)
            
            if chunk.strip():
                chunks.append({
                    "content": chunk.strip(),
                    "chunk_id": len(chunks)
                })
            
            # 移动到下一个块，考虑重叠
            start = end - self.chunk_overlap
            
            # 防止无限循环
            if start >= len(text) or len(chunks) == 0:
                break
        
        return chunks
